// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <nebula/array/builder_dict.h>

#include <cstdint>
#include <utility>

#include <nebula/array/dict_internal.h>
#include <turbo/utility/status.h>
#include <nebula/types/type.h>
#include <nebula/types/type_traits.h>
#include <turbo/base/checked_cast.h>
#include <nebula/util/hashing.h>
#include <turbo/log/logging.h>
#include <nebula/core/visit_type_inline.h>

namespace nebula::internal {

    class DictionaryMemoTable::DictionaryMemoTableImpl {
        // Type-dependent visitor for memo table initialization
        struct MemoTableInitializer {
            std::shared_ptr<DataType> value_type_;
            MemoryPool *pool_;
            std::unique_ptr<MemoTable> *memo_table_;

            template<typename T>
            enable_if_no_memoize<T, turbo::Status> Visit(const T &) {
                return turbo::unimplemented_error("Initialization of ", value_type_->to_string(),
                                                  " memo table is not implemented");
            }

            template<typename T>
            enable_if_memoize<T, turbo::Status> Visit(const T &) {
                using MemoTable = typename DictionaryTraits<T>::MemoTableType;
                memo_table_->reset(new MemoTable(pool_, 0));
                return turbo::OkStatus();
            }
        };

        // Type-dependent visitor for memo table insertion
        struct ArrayValuesInserter {
            DictionaryMemoTableImpl *impl_;
            const Array &values_;

            template<typename T>
            turbo::Status Visit(const T &type) {
                using ArrayType = typename TypeTraits<T>::ArrayType;
                return InsertValues(type, turbo::checked_cast<const ArrayType &>(values_));
            }

        private:
            template<typename T, typename ArrayType>
            enable_if_no_memoize<T, turbo::Status> InsertValues(const T &type, const ArrayType &) {
                return turbo::unimplemented_error("Inserting array values of ", type,
                                                  " is not implemented");
            }

            template<typename T, typename ArrayType>
            enable_if_memoize<T, turbo::Status> InsertValues(const T &, const ArrayType &array) {
                if (array.null_count() > 0) {
                    return turbo::invalid_argument_error("Cannot insert dictionary values containing nulls");
                }
                for (int64_t i = 0; i < array.length(); ++i) {
                    int32_t unused_memo_index;
                    TURBO_RETURN_NOT_OK(impl_->GetOrInsert<T>(array.get_view(i), &unused_memo_index));
                }
                return turbo::OkStatus();
            }
        };

        // Type-dependent visitor for building ArrayData from memo table
        struct ArrayDataGetter {
            std::shared_ptr<DataType> value_type_;
            MemoTable *memo_table_;
            MemoryPool *pool_;
            int64_t start_offset_;
            std::shared_ptr<ArrayData> *out_;

            template<typename T>
            enable_if_no_memoize<T, turbo::Status> Visit(const T &) {
                return turbo::unimplemented_error("Getting array data of ", value_type_,
                                                  " is not implemented");
            }

            template<typename T>
            enable_if_memoize<T, turbo::Status> Visit(const T &) {
                using ConcreteMemoTable = typename DictionaryTraits<T>::MemoTableType;
                auto memo_table = turbo::checked_cast<ConcreteMemoTable *>(memo_table_);
                TURBO_MOVE_OR_RAISE(*out_, DictionaryTraits<T>::GetDictionaryArrayData(
                        pool_, value_type_, *memo_table, start_offset_));
                return turbo::OkStatus();
            }
        };

    public:
        DictionaryMemoTableImpl(MemoryPool *pool, std::shared_ptr<DataType> type)
                : pool_(pool), type_(std::move(type)), memo_table_(nullptr) {
            MemoTableInitializer visitor{type_, pool_, &memo_table_};
            KCHECK_OK(visit_type_inline(*type_, &visitor));
        }

        turbo::Status InsertValues(const Array &array) {
            if (!array.type()->equals(*type_)) {
                return turbo::invalid_argument_error("Array value type does not match memo type: ",
                                                     array.type()->to_string());
            }
            ArrayValuesInserter visitor{this, array};
            return visit_type_inline(*array.type(), &visitor);
        }

        template<typename PhysicalType,
                typename CType = typename DictionaryValue<PhysicalType>::type>
        turbo::Status GetOrInsert(CType value, int32_t *out) {
            using ConcreteMemoTable = typename DictionaryTraits<PhysicalType>::MemoTableType;
            return turbo::checked_cast<ConcreteMemoTable *>(memo_table_.get())->GetOrInsert(value, out);
        }

        turbo::Status GetArrayData(int64_t start_offset, std::shared_ptr<ArrayData> *out) {
            ArrayDataGetter visitor{type_, memo_table_.get(), pool_, start_offset, out};
            return visit_type_inline(*type_, &visitor);
        }

        int32_t size() const { return memo_table_->size(); }

    private:
        MemoryPool *pool_;
        std::shared_ptr<DataType> type_;
        std::unique_ptr<MemoTable> memo_table_;
    };

    DictionaryMemoTable::DictionaryMemoTable(MemoryPool *pool,
                                             const std::shared_ptr<DataType> &type)
            : impl_(new DictionaryMemoTableImpl(pool, type)) {}

    DictionaryMemoTable::DictionaryMemoTable(MemoryPool *pool,
                                             const std::shared_ptr<Array> &dictionary)
            : impl_(new DictionaryMemoTableImpl(pool, dictionary->type())) {
        KCHECK_OK(impl_->InsertValues(*dictionary));
    }

    DictionaryMemoTable::~DictionaryMemoTable() = default;

#define GET_OR_INSERT(NEBULA_TYPE)                                           \
      turbo::Status DictionaryMemoTable::GetOrInsert(                                  \
          const NEBULA_TYPE*, typename NEBULA_TYPE::c_type value, int32_t* out) { \
        return impl_->GetOrInsert<NEBULA_TYPE>(value, out);                      \
      }

    GET_OR_INSERT(BooleanType)

    GET_OR_INSERT(Int8Type)

    GET_OR_INSERT(Int16Type)

    GET_OR_INSERT(Int32Type)

    GET_OR_INSERT(Int64Type)

    GET_OR_INSERT(UInt8Type)

    GET_OR_INSERT(UInt16Type)

    GET_OR_INSERT(UInt32Type)

    GET_OR_INSERT(UInt64Type)

    GET_OR_INSERT(Fp32Type)

    GET_OR_INSERT(Fp64Type)

    GET_OR_INSERT(DurationType);

    GET_OR_INSERT(TimestampType);

    GET_OR_INSERT(Date32Type);

    GET_OR_INSERT(Date64Type);

    GET_OR_INSERT(Time32Type);

    GET_OR_INSERT(Time64Type);

    GET_OR_INSERT(MonthDayNanoIntervalType);

    GET_OR_INSERT(DayTimeIntervalType);

    GET_OR_INSERT(MonthIntervalType);

#undef GET_OR_INSERT

    turbo::Status DictionaryMemoTable::GetOrInsert(const BinaryType *, std::string_view value,
                                                   int32_t *out) {
        return impl_->GetOrInsert<BinaryType>(value, out);
    }

    turbo::Status DictionaryMemoTable::GetOrInsert(const BinaryViewType *, std::string_view value,
                                                   int32_t *out) {
        // create BinaryArray dictionary for now
        return impl_->GetOrInsert<BinaryType>(value, out);
    }

    turbo::Status DictionaryMemoTable::GetOrInsert(const LargeBinaryType *, std::string_view value,
                                                   int32_t *out) {
        return impl_->GetOrInsert<LargeBinaryType>(value, out);
    }

    turbo::Status DictionaryMemoTable::GetArrayData(int64_t start_offset,
                                                    std::shared_ptr<ArrayData> *out) {
        return impl_->GetArrayData(start_offset, out);
    }

    turbo::Status DictionaryMemoTable::InsertValues(const Array &array) {
        return impl_->InsertValues(array);
    }

    int32_t DictionaryMemoTable::size() const { return impl_->size(); }

}  // namespace nebula::internal
